import pymysql
import pandas as pd
class MysqlUtils(object):
    def __init__(self):
        self.conn = pymysql.connect(
            host='127.0.0.1',
            user="root",
            password="root",
            db="tusharel",
            port=3306,
            charset="utf8"
        )
class classIfication(object):
    #获取财务相关数据
    def get_fina_indicator(self, conn):
        cursor = conn.cursor(cursor=pymysql.cursors.DictCursor)

        sql = """SELECT ts_code, ann_date, eps, total_revenue_ps, undist_profit_ps,
        gross_margin, fcff, fcfe, tangible_asset, bps, grossprofit_margin, npta, roic
        FROM financial_data WHERE ann_date >= '2023-01-01' and ann_date<'2024-01-01'"""

        cursor.execute(sql)
        ret = cursor.fetchall()
        #print(ret)
        df = pd.DataFrame(ret)
        df1 = df.dropna(subset=['eps', 'total_revenue_ps', 'undist_profit_ps',
                                'gross_margin','fcff', 'fcfe','tangible_asset',
                                 'bps', 'grossprofit_margin', 'npta', 'roic'])
        #重建索引
        df1 = df1.reset_index(drop=False)
        print(df1.head)

        return df1

        

    #获取日线行情
    def get_daily(self, conn, df1):
        cursor = conn.cursor(cursor=pymysql.cursor.DictCursor)
        new_list = []
        for i in range(len(df1['ts_code'])):
                try:
                    print(i)
                    ann_date_str = df1['ann_date'][i].strftime('%Y%m%d')
                    sql = 'select trade_date, closes, opens, high, low from date_l where ts_code = \''\
                         + df1['ts_cide'][i] + '\' and trade_date > Date(\''+ ann_date_str \
                             +'\') order by trade_date limit 20'
                    cursor.execute(sql)
                    ret = cursor.fetchall()

                    df2 = pd.DataFrame(ret)
                    if len(df2) > 0:
                        max_close = df2['closes'].max()
                        min_close = df2['closes'].min()
                        the_close = df2['closes'].iloc[1]
                        new_list.append({
                            'ts_code': df1['ts_code'][i],
                            'ann_date': df1['ann_date'][i],
                            'max_close': max_close,
                            'min_close': min_close,
                            'the_close': the_close,
                            'eps': df1['eps'][i],
                            'total_revenue_ps': df1['total_revenue_ps'][i],
                            'undist_profit_ps': df1['undist_profit_ps'][i],
                            'gross_margin': df1['gross_margin'][i],
                            'fcff': df1['fcff'][i], 
                            'fcfe': df1['fcfe'][i],
                            'tangible_asset': df1['tangible_asset'][i], 
                            'bps': df1['bps'][i], 
                            'grossprofit_margin': df1['grossprofit_margin'][i], 
                            'npta': df1['npta'][i], 
                            'roic': df1['roic'][i],
                        })
                except Exception as e:
                    print(e)
        df3 = pd.DataFrame(new_list)
        df3.to_csv('daily.csv', index=False)

if __name__ == '__main__':
    mu = MysqlUtils()
    ci = classIfication()
    ci.get_fina_indicator(mu.conn)